{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "#### CNN实现MNIST分类--pytorch version\n",
    "\n",
    "> python 3.7\n",
    ">\n",
    "> 2023/02/13"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import os\n",
    "import time"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "data": {
      "text/plain": "<torch._C.Generator at 0x106584430>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1) # 为torch中的随机操作设计一个随机种子，如此每次结果都一样"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "超参\n",
    "\"\"\"\n",
    "batch_size = 128\n",
    "epochs = 4\n",
    "lr = 0.01"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "download_mnist = False\n",
    "data_dir = './../dataset/'\n",
    "if os.path.exists(data_dir):\n",
    "    download_mnist = False\n",
    "else:\n",
    "    download_mnist = True"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data_size 60000\n",
      "batch_size 128\n"
     ]
    }
   ],
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root=data_dir,\n",
    "                                           train=True,\n",
    "                                           transform=transforms.ToTensor(),\n",
    "                                           download=download_mnist)\n",
    "\n",
    "test_dataset = torchvision.datasets.MNIST(root=data_dir,\n",
    "                                          train=False,\n",
    "                                          transform=transforms.ToTensor())\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset,\n",
    "                                          batch_size=batch_size,\n",
    "                                          shuffle=False)\n",
    "print('data_size',len(train_dataset))\n",
    "print('batch_size',batch_size)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNN,self).__init__()\n",
    "        # input shape (1,28,28)\n",
    "        self.conv1 = nn.Sequential(\n",
    "            # input shape (1,28,28)\n",
    "            nn.Conv2d(\n",
    "                in_channels=1, # 深度 只有一个通道\n",
    "                out_channels=16, # filter卷积核个数 输出通道数\n",
    "                kernel_size=3, # 卷积核大小 3*3\n",
    "                stride=1, # 步长\n",
    "                padding=1 # 填充\n",
    "            ),\n",
    "            # input shape (16,28,28)\n",
    "            nn.ReLU(), # activate\n",
    "            nn.MaxPool2d(kernel_size=2), # 2*2的最大汇聚层\n",
    "        )\n",
    "        # output shape (16,14,14)\n",
    "\n",
    "        # input shape (16,14,14)\n",
    "        self.conv2 = nn.Sequential(\n",
    "            # input shape (16,14,14)\n",
    "            nn.Conv2d(\n",
    "                in_channels=16,\n",
    "                out_channels=32,\n",
    "                kernel_size=3,\n",
    "                stride=1,\n",
    "                padding=1\n",
    "            ),\n",
    "            # input shape (32,14,14)\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "        )\n",
    "        # output shape (32,7,7)\n",
    "        self.out = nn.Sequential(\n",
    "            nn.Linear(32*7*7,10), # 全连接层\n",
    "            nn.Softmax(dim=-1) # 激活函数\n",
    "        )\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = x.view(x.size(0),-1)\n",
    "        x = self.out(x)\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CNN(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (out): Sequential(\n",
      "    (0): Linear(in_features=1568, out_features=10, bias=True)\n",
      "    (1): Softmax(dim=-1)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "net = CNN().to(device)\n",
    "print(net)\n",
    "\n",
    "optimizer = torch.optim.Adam(net.parameters(),lr=lr)\n",
    "criterion = nn.CrossEntropyLoss()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total_step 469\n",
      "Epoch [1/4], Step [100/469], Loss: 1.4972\n",
      "Epoch [1/4], Step [200/469], Loss: 1.4941\n",
      "Epoch [1/4], Step [300/469], Loss: 1.5057\n",
      "Epoch [1/4], Step [400/469], Loss: 1.4886\n",
      "one epoch cost 16.80 seconds\n",
      "Epoch [2/4], Step [100/469], Loss: 1.4874\n",
      "Epoch [2/4], Step [200/469], Loss: 1.4956\n",
      "Epoch [2/4], Step [300/469], Loss: 1.4802\n",
      "Epoch [2/4], Step [400/469], Loss: 1.4679\n",
      "one epoch cost 15.99 seconds\n",
      "Epoch [3/4], Step [100/469], Loss: 1.4832\n",
      "Epoch [3/4], Step [200/469], Loss: 1.5175\n",
      "Epoch [3/4], Step [300/469], Loss: 1.4739\n",
      "Epoch [3/4], Step [400/469], Loss: 1.4873\n",
      "one epoch cost 16.09 seconds\n",
      "Epoch [4/4], Step [100/469], Loss: 1.4687\n",
      "Epoch [4/4], Step [200/469], Loss: 1.5082\n",
      "Epoch [4/4], Step [300/469], Loss: 1.4690\n",
      "Epoch [4/4], Step [400/469], Loss: 1.4734\n",
      "one epoch cost 16.16 seconds\n"
     ]
    }
   ],
   "source": [
    "total_step = len(train_loader)\n",
    "print('total_step {}'.format(total_step))\n",
    "for epoch in range(epochs):\n",
    "    ts = time.time()\n",
    "    for i,(images,labels) in enumerate(train_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        outputs = net(images)\n",
    "        loss = criterion(outputs,labels)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if (i+1) % 100 == 0:\n",
    "            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'\n",
    "                   .format(epoch+1, epochs, i+1, total_step, loss.item()))\n",
    "    te = time.time()\n",
    "    print('one epoch cost {:.2f} seconds'.format(te-ts))\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy of the model on the 10000 test images: 98.04 %\n"
     ]
    }
   ],
   "source": [
    "# Test the model\n",
    "net.eval()\n",
    "with torch.no_grad():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = net(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
